import json
import os

import numpy as np
import torch
import torch.nn as nn
import vedo
from vedo import Mesh
from scipy.spatial import distance_matrix

import utils
from meshsegnet import MeshSegNet
from meshclsnet import MeshClsNet

os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"

if __name__ == '__main__':
    gpu_id = utils.get_avail_gpu()
    torch.cuda.set_device(gpu_id)

    with open('config.json', 'r') as f:
        config = json.load(f)

    mod = config["mod"]

    i_mesh = config["data"]["dataset_path"] + config["data"]["data_path"] + mod + "/upper4.ply"

    mesh = Mesh(i_mesh)
    colors = mesh.cellcolors[:, 0:3] / 255.
    points = mesh.points()
    mean_cell_centers = mesh.center_of_mass()
    points[:, 0:3] -= mean_cell_centers[0:3]

    ids = np.array(mesh.faces())
    cells = points[ids].reshape(mesh.ncells, 9).astype(dtype='float32')

    mesh.compute_normals()
    normals = mesh.celldata['Normals']

    barycenters = mesh.cell_centers()
    barycenters -= mean_cell_centers[0:3]

    maxs = points.max(axis=0)
    mins = points.min(axis=0)
    means = points.mean(axis=0)
    stds = points.std(axis=0)
    nmeans = normals.mean(axis=0)
    nstds = normals.std(axis=0)

    for i in range(3):
        cells[:, i] = (cells[:, i] - means[i]) / stds[i]
        cells[:, i + 3] = (cells[:, i + 3] - means[i]) / stds[i]
        cells[:, i + 6] = (cells[:, i + 6] - means[i]) / stds[i]
        barycenters[:, i] = (barycenters[:, i] - mins[i]) / (maxs[i] - mins[i])
        normals[:, i] = (normals[:, i] - nmeans[i]) / nstds[i]

    X = np.column_stack((cells, barycenters, normals, colors))

    A_S = np.zeros([X.shape[0], X.shape[0]], dtype='float32')
    A_L = np.zeros([X.shape[0], X.shape[0]], dtype='float32')
    D = distance_matrix(X[:, 9:12], X[:, 9:12])
    A_S[D < 0.1] = 1.0
    A_S = A_S / np.dot(np.sum(A_S, axis=1, keepdims=True), np.ones((1, X.shape[0])))

    A_L[D < 0.2] = 1.0
    A_L = A_L / np.dot(np.sum(A_L, axis=1, keepdims=True), np.ones((1, X.shape[0])))

    models_seg_path = config["data"]["models_seg_path"]
    models_seg_name = config["data"]["models_seg_name"] + '{}.tar'.format(mod)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model_seg = MeshSegNet(num_classes=15, num_channels=15).to(device, dtype=torch.float)

    checkpoint = torch.load(os.path.join(models_seg_path, models_seg_name), map_location='cpu')
    model_seg.load_state_dict(checkpoint['model_state_dict'])

    torch.backends.cudnn.benchmark = True
    torch.backends.cudnn.enabled = True

    model_seg.eval()

    X_seg = torch.from_numpy(X[:, 0:15]).to(device, dtype=torch.float)
    X_seg = X_seg.transpose(1, 0)
    X_seg = X_seg.reshape([-1, X_seg.shape[0], X_seg.shape[1]])
    A_S = A_S.reshape([1, A_S.shape[0], A_S.shape[1]])
    A_L = A_L.reshape([1, A_L.shape[0], A_L.shape[1]])
    A_S = torch.from_numpy(A_S).to(device, dtype=torch.float)
    A_L = torch.from_numpy(A_L).to(device, dtype=torch.float)

    with torch.no_grad():
        seg_output = model_seg(X_seg, A_S, A_L).to(device, dtype=torch.float)
        seg_matrix = seg_output.cpu().numpy()

    X_pre = np.zeros([6000, X.shape[1]], dtype='float32')

    valid_idx = np.array([i for i in range(X.shape[0])])
    selected_idx = np.sort(np.random.choice(valid_idx, size=6000, replace=False))

    seg_matrix = seg_matrix.reshape(seg_matrix.shape[1], seg_matrix.shape[2])

    c_seg_matrix = np.zeros_like(seg_matrix)
    c_seg_matrix[np.arange(len(seg_matrix)), seg_matrix.argmax(1)] = 1
    value_index = c_seg_matrix.argmax(1)

    centers = mesh.cell_centers()
    area_centers = np.zeros([15, 3])
    area_sum = [0 for i in range(15)]

    for i in range(len(value_index)):
        index = value_index[i]
        center = centers[i]
        # print(index, center)
        if index != 0:
            area_sum[index] += 1
            area_centers[index] = area_centers[index] + center

    for j in range(1, 15):
        if area_sum[j] != 0:
            area_centers[j, 0] = area_centers[j, 0] / area_sum[j]
            area_centers[j, 1] = area_centers[j, 1] / area_sum[j]
            area_centers[j, 2] = area_centers[j, 2] / area_sum[j]

    for i in range(len(value_index)):
        index = value_index[i]

        if index == 0:
            center = centers[i]
            distance = 10000
            nearest = 0
            for j in range(1, 15):
                if area_centers[j][0] != 0.:
                    _distance = np.sqrt(np.sum((area_centers[j] - center) ** 2))
                    if _distance < distance:
                        distance = _distance
                        nearest = j
            c_seg_matrix[i, nearest] = 1

    seg_matrix = c_seg_matrix


    X_pre[:] = X[selected_idx, :]

    seg_matrix = seg_matrix[selected_idx, :]

    seg_matrix = torch.from_numpy(seg_matrix.reshape(1, seg_matrix.shape[0], seg_matrix.shape[1])).to(device)

    X_pre = X_pre.transpose(1, 0)

    X_pre = torch.from_numpy(X_pre.reshape(1, X_pre.shape[0], X_pre.shape[1])).to(device)

    model_path = config["train"]["model_path"]
    model_name = '{}{}_best.tar'.format(config["train"]["model_name"], mod)

    model = MeshClsNet(num_classes=5, num_channels=18).to(device, dtype=torch.float)

    checkpoint = torch.load(os.path.join(model_path + mod, model_name), map_location='cpu')
    model.load_state_dict(checkpoint['model_state_dict'])

    model.eval()
    with torch.no_grad():
        output = model(X_pre, seg_matrix).to(device, dtype=torch.float)
        output = output.cpu()
        output = output.reshape(output.shape[1], output.shape[2]).transpose(0, 1)
        print(output.argmax(axis=1))

